from keras.layers import Input, Conv2D, MaxPool2D, Dense, Activation, Flatten
from keras.models import Model
from keras.optimizers import adam_v2
from keras.callbacks import ReduceLROnPlateau, ModelCheckpoint, TensorBoard
import os
import numpy as np


def model():

    _in = Input(shape=(32,32,3))                    #输入层，形状为（32，32，3）

    x = Conv2D(32, (3,3), padding='same')(_in)      #卷积层 过滤器（filters）数量=32,kernel_size=(3,3),边缘填充（padding）-->(x,x,32)
    x = MaxPool2D((2,2))(x)                         #池化层 最大池化(2*2)   -->(x/2,x/2,32)
    x = Activation("relu")(x)                       #激活函数 x=max(0,x) 引入激活函数是为了拟合非线性问题否则神经网络没用，使用relu：计算量小，不会梯度消失，稀疏防止过拟合

    x = Conv2D(64, (3,3), padding='same')(x)        #-->(x,x,64)
    x = MaxPool2D((2,2))(x)                         #-->(x/2,x/2,64)
    x = Activation("relu")(x)                       #-->(x/2,x/2,64)

    x = Conv2D(128, (3,3), padding='same')(x)       #-->(x,x,128)
    x = MaxPool2D((2,2))(x)                         #-->(x/2,x/2,128)
    x = Activation("relu")(x)                       #-->(x/2,x/2,128)

    x = Flatten()(x)                                #展平函数将数组展平为1维，卷积层与全链接层过度
    x = Dense(10)(x)                                #最后接全连接层，类别为10，所以共10个单元
    x = Activation("softmax")(x)                    #得分差异化，概率归一化，配合多分类交叉熵loss function

    return Model(_in,x)


if __name__ == "__main__":

    data_train = np.load('./data_train.npy')    #载入文件
    label_train = np.load('./label_train.npy')
    data_test = np.load('./data_test.npy')
    label_test = np.load('./label_test.npy')

    if not (os.path.exists('./models')):        #创建模型文件夹
        os.mkdir("./models")

    model = model()             #初始化网络模型
    model.summary()             #模型的参数输出


    #优化器，随机梯度下降，自适应更改学习率
    # #配置优化器，损失函数（多分类交叉熵），评价参数，准确率
    model.compile(optimizer=adam_v2.Adam(lr=0.001), loss='categorical_crossentropy', metrics=["acc"])


    # 15轮参数不优化就减小学习率，步长变小
    # 保存模型
    # log可视化
    reduce_lr = ReduceLROnPlateau(patience=15)
    save_weights = ModelCheckpoint("./models/model_{epoch:02d}_{val_acc:.4f}.h5", save_best_only=True, monitor='val_acc')
    log=TensorBoard(log_dir="logs")
    callbacks = [save_weights, reduce_lr, log]

    # 开始迭代优化
    model.fit(data_train, label_train, validation_data=(data_test, label_test), epochs = 50, batch_size=32, callbacks=callbacks)   #开始训练，训练100轮，一次喂进去32个数据，每轮结束后用test评估，然后返回评估结果